from flask import Flask, jsonify, request
import cv2
import numpy as np
import base64
import requests
import time
import threading
import json
from datetime import datetime
import os

app = Flask(__name__)

class ReconstructionService:
    def __init__(self, video_server_url="http://host.docker.internal:5000"):
        self.video_server_url = video_server_url
        self.is_processing = False
        self.current_frame = None
        self.frame_buffer = []
        self.max_buffer_size = 10
        self.processing_stats = {
            'frames_processed': 0,
            'start_time': None,
            'current_fps': 0
        }
        
    def get_frame_from_server(self):
        """从视频流服务器获取当前帧"""
        try:
            response = requests.get(f"{self.video_server_url}/frame", timeout=1.0)
            if response.status_code == 200:
                data = response.json()
                if data['status'] == 'success':
                    # 解码 base64 图像
                    frame_data = base64.b64decode(data['frame'])
                    frame_array = np.frombuffer(frame_data, dtype=np.uint8)
                    frame = cv2.imdecode(frame_array, cv2.IMREAD_COLOR)
                    return frame, data.get('timestamp', time.time())
            return None, None
        except Exception as e:
            print(f"获取帧失败: {e}")
            return None, None
    
    def get_server_status(self):
        """获取视频流服务器状态"""
        try:
            response = requests.get(f"{self.video_server_url}/status", timeout=1.0)
            if response.status_code == 200:
                return response.json()
            return None
        except Exception as e:
            print(f"获取服务器状态失败: {e}")
            return None
    
    def start_video_stream(self):
        """启动视频流"""
        try:
            response = requests.get(f"{self.video_server_url}/start", timeout=2.0)
            if response.status_code == 200:
                data = response.json()
                if data['status'] == 'success':
                    print("视频流已启动")
                    return True
            return False
        except Exception as e:
            print(f"启动视频流失败: {e}")
            return False
    
    def process_frame(self, frame):
        """处理单帧图像（这里可以添加你的三维重建算法）"""
        if frame is None:
            return None
        
        # 这里添加你的三维重建处理逻辑
        # 示例：简单的图像预处理
        processed_frame = frame.copy()
        
        # 转换为灰度图
        gray = cv2.cvtColor(processed_frame, cv2.COLOR_BGR2GRAY)
        
        # 边缘检测
        edges = cv2.Canny(gray, 50, 150)
        
        # 查找轮廓
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 绘制轮廓
        cv2.drawContours(processed_frame, contours, -1, (0, 255, 0), 2)
        
        # 添加处理信息
        cv2.putText(processed_frame, f"Contours: {len(contours)}", 
                   (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        
        return processed_frame
    
    def start_processing(self):
        """开始处理视频流"""
        if self.is_processing:
            return False
        
        # 启动视频流
        if not self.start_video_stream():
            return False
        
        self.is_processing = True
        self.processing_stats['start_time'] = time.time()
        self.processing_stats['frames_processed'] = 0
        
        def processing_loop():
            while self.is_processing:
                frame, timestamp = self.get_frame_from_server()
                if frame is not None:
                    self.current_frame = frame
                    
                    # 处理帧
                    processed_frame = self.process_frame(frame)
                    
                    # 添加到缓冲区
                    self.frame_buffer.append({
                        'original': frame,
                        'processed': processed_frame,
                        'timestamp': timestamp
                    })
                    
                    # 限制缓冲区大小
                    if len(self.frame_buffer) > self.max_buffer_size:
                        self.frame_buffer.pop(0)
                    
                    # 更新统计信息
                    self.processing_stats['frames_processed'] += 1
                    elapsed_time = time.time() - self.processing_stats['start_time']
                    if elapsed_time > 0:
                        self.processing_stats['current_fps'] = self.processing_stats['frames_processed'] / elapsed_time
                
                time.sleep(0.033)  # ~30 FPS
        
        threading.Thread(target=processing_loop, daemon=True).start()
        return True
    
    def stop_processing(self):
        """停止处理"""
        self.is_processing = False
    
    def get_processing_stats(self):
        """获取处理统计信息"""
        return self.processing_stats.copy()
    
    def get_latest_frame(self):
        """获取最新处理的帧"""
        if self.frame_buffer:
            latest = self.frame_buffer[-1]
            # 将处理后的帧编码为 base64
            _, buffer = cv2.imencode('.jpg', latest['processed'])
            frame_base64 = base64.b64encode(buffer).decode('utf-8')
            return {
                'frame': frame_base64,
                'timestamp': latest['timestamp'],
                'contours_count': len(cv2.findContours(
                    cv2.Canny(cv2.cvtColor(latest['original'], cv2.COLOR_BGR2GRAY), 50, 150),
                    cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
                )[0])
            }
        return None

# 全局重建服务实例
reconstruction_service = ReconstructionService()

@app.route('/')
def index():
    return jsonify({
        'status': '3D Reconstruction Service Running',
        'endpoints': {
            '/start': 'Start reconstruction processing',
            '/stop': 'Stop reconstruction processing',
            '/status': 'Get processing status',
            '/frame': 'Get latest processed frame',
            '/stats': 'Get processing statistics'
        }
    })

@app.route('/start')
def start_reconstruction():
    """启动三维重建处理"""
    if reconstruction_service.start_processing():
        return jsonify({'status': 'success', 'message': '三维重建处理已启动'})
    else:
        return jsonify({'status': 'error', 'message': '无法启动处理'}), 500

@app.route('/stop')
def stop_reconstruction():
    """停止三维重建处理"""
    reconstruction_service.stop_processing()
    return jsonify({'status': 'success', 'message': '三维重建处理已停止'})

@app.route('/status')
def get_status():
    """获取处理状态"""
    video_status = reconstruction_service.get_server_status()
    return jsonify({
        'status': 'success',
        'is_processing': reconstruction_service.is_processing,
        'video_server_status': video_status,
        'frame_buffer_size': len(reconstruction_service.frame_buffer)
    })

@app.route('/frame')
def get_latest_frame():
    """获取最新处理的帧"""
    frame_data = reconstruction_service.get_latest_frame()
    if frame_data:
        return jsonify({
            'status': 'success',
            'frame': frame_data
        })
    else:
        return jsonify({'status': 'error', 'message': '无可用帧'}), 404

@app.route('/stats')
def get_stats():
    """获取处理统计信息"""
    stats = reconstruction_service.get_processing_stats()
    return jsonify({
        'status': 'success',
        'stats': stats
    })

if __name__ == '__main__':
    print("启动三维重建服务...")
    print("访问 http://localhost:8000 查看API文档")
    print("确保视频流服务器在 http://host.docker.internal:5000 运行")
    app.run(host='0.0.0.0', port=8000, debug=False, threaded=True) 